{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 14\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import functions from the modsim.py module\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code from previous chapters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_system(beta, gamma):\n",
    "    \"\"\"Make a system object for the SIR model.\n",
    "    \n",
    "    beta: contact rate in days\n",
    "    gamma: recovery rate in days\n",
    "    \n",
    "    returns: System object\n",
    "    \"\"\"\n",
    "    init = State(S=89, I=1, R=0)\n",
    "    init /= np.sum(init)\n",
    "\n",
    "    t0 = 0\n",
    "    t_end = 7 * 14\n",
    "\n",
    "    return System(init=init, t0=t0, t_end=t_end,\n",
    "                  beta=beta, gamma=gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func(state, t, system):\n",
    "    \"\"\"Update the SIR model.\n",
    "    \n",
    "    state: State (s, i, r)\n",
    "    t: time\n",
    "    system: System object\n",
    "    \n",
    "    returns: State (sir)\n",
    "    \"\"\"\n",
    "    s, i, r = state\n",
    "\n",
    "    infected = system.beta * i * s    \n",
    "    recovered = system.gamma * i\n",
    "    \n",
    "    s -= infected\n",
    "    i += infected - recovered\n",
    "    r += recovered\n",
    "    \n",
    "    return State(S=s, I=i, R=r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Runs a simulation of the system.\n",
    "        \n",
    "    system: System object\n",
    "    update_func: function that updates state\n",
    "    \n",
    "    returns: TimeFrame\n",
    "    \"\"\"\n",
    "    init, t0, t_end = system.init, system.t0, system.t_end\n",
    "    \n",
    "    frame = TimeFrame(columns=init.index)\n",
    "    frame.row[t0] = init\n",
    "    \n",
    "    for t in linrange(t0, t_end):\n",
    "        frame.row[t+1] = update_func(frame.row[t], t, system)\n",
    "    \n",
    "    return frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_total_infected(results):\n",
    "    \"\"\"Fraction of population infected during the simulation.\n",
    "    \n",
    "    results: DataFrame with columns S, I, R\n",
    "    \n",
    "    returns: fraction of population\n",
    "    \"\"\"\n",
    "    return get_first_value(results.S) - get_last_value(results.S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_beta(beta_array, gamma):\n",
    "    \"\"\"Sweep a range of values for beta.\n",
    "    \n",
    "    beta_array: array of beta values\n",
    "    gamma: recovery rate\n",
    "    \n",
    "    returns: SweepSeries that maps from beta to total infected\n",
    "    \"\"\"\n",
    "    sweep = SweepSeries()\n",
    "    for beta in beta_array:\n",
    "        system = make_system(beta, gamma)\n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[system.beta] = calc_total_infected(results)\n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_parameters(beta_array, gamma_array):\n",
    "    \"\"\"Sweep a range of values for beta and gamma.\n",
    "    \n",
    "    beta_array: array of infection rates\n",
    "    gamma_array: array of recovery rates\n",
    "    \n",
    "    returns: SweepFrame with one row for each beta\n",
    "             and one column for each gamma\n",
    "    \"\"\"\n",
    "    frame = SweepFrame(columns=gamma_array)\n",
    "    for gamma in gamma_array:\n",
    "        frame[gamma] = sweep_beta(beta_array, gamma)\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Contact number"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the `SweepFrame` from the previous chapter, with one row for each value of `beta` and one column for each value of `gamma`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_array = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0 , 1.1]\n",
    "gamma_array = [0.2, 0.4, 0.6, 0.8]\n",
    "frame = sweep_parameters(beta_array, gamma_array)\n",
    "frame.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following loop shows how we can loop through the columns and rows of the `SweepFrame`.  With 11 rows and 4 columns, there are 44 elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in frame.columns:\n",
    "    column = frame[gamma]\n",
    "    for beta in column.index:\n",
    "        frac_infected = column[beta]\n",
    "        print(beta, gamma, frac_infected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can wrap that loop in a function and plot the results.  For each element of the `SweepFrame`, we have `beta`, `gamma`, and `frac_infected`, and we plot `beta/gamma` on the x-axis and `frac_infected` on the y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sweep_frame(frame):\n",
    "    \"\"\"Plot the values from a SweepFrame.\n",
    "    \n",
    "    For each (beta, gamma), compute the contact number,\n",
    "    beta/gamma\n",
    "    \n",
    "    frame: SweepFrame with one row per beta, one column per gamma\n",
    "    \"\"\"\n",
    "    for gamma in frame.columns:\n",
    "        column = frame[gamma]\n",
    "        for beta in column.index:\n",
    "            frac_infected = column[beta]\n",
    "            plot(beta/gamma, frac_infected, 'ro')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_sweep_frame(frame)\n",
    "\n",
    "decorate(xlabel='Contact number (beta/gamma)',\n",
    "         ylabel='Fraction infected')\n",
    "\n",
    "savefig('figs/chap14-fig01.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that the ratio `beta/gamma`, called the \"contact number\" is sufficient to predict the total number of infections; we don't have to know `beta` and `gamma` separately.\n",
    "\n",
    "We can see that in the previous plot: when we plot the fraction infected versus the contact number, the results fall close to a curve."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the book we figured out the relationship between $c$ and $s_{\\infty}$ analytically.  Now we can compute it for a range of values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_inf_array = linspace(0.0001, 0.9999, 101);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_array = log(s_inf_array) / (s_inf_array - 1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`total_infected` is the change in $s$ from the beginning to the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "frac_infected = 1 - s_inf_array\n",
    "frac_infected_series = Series(frac_infected, index=c_array);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot the analytic results and compare them to the simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_sweep_frame(frame)\n",
    "plot(frac_infected_series, label='Analysis')\n",
    "\n",
    "decorate(xlabel='Contact number (c)',\n",
    "         ylabel='Fraction infected')\n",
    "\n",
    "savefig('figs/chap14-fig02.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The agreement is generally good, except for values of `c` less than 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**  If we didn't know about contact numbers, we might have explored other possibilities, like the difference between `beta` and `gamma`, rather than their ratio.\n",
    "\n",
    "Write a version of `plot_sweep_frame`, called `plot_sweep_frame_difference`, that plots the fraction infected versus the difference `beta-gamma`.\n",
    "\n",
    "What do the results look like, and what does that imply? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Suppose you run a survey at the end of the semester and find that 26% of students had the Freshman Plague at some point.\n",
    "\n",
    "What is your best estimate of `c`?\n",
    "\n",
    "Hint: if you print `frac_infected_series`, you can read off the answer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
